package com.kongxiang.cubebit.llm.service;

import com.kongxiang.cubebit.base.common.dto.R;
import com.kongxiang.cubebit.base.common.exceptions.BaseException;
import com.kongxiang.cubebit.llm.dto.ReRankRequest;
import com.kongxiang.cubebit.llm.dto.ReRankResponse;
import com.kongxiang.cubebit.vectordb.services.VectorStoreService;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Service
public class RagService {

    @Autowired
    private VectorStoreService vectorStoreService;

    @Autowired
    private OpenAIService openAIService;

    /**
     * 混合搜索功能，结合向量搜索和关键字搜索
     *
     * @param query 查询条件
     * @param k 返回结果数量
     * @param threshold 相似度阈值
     * @return 搜索结果
     */
    public List<Document> ragSearch(String query, int k, double threshold)  {
        // 向量搜索
        List<Document> vectorResults = vectorStoreService.search(query, k, threshold);
        // vectorResult非空检验和报错
        if (vectorResults == null || vectorResults.isEmpty()) {
            return Collections.emptyList();
        }
        List<String> rerankDocuments = new ArrayList<>();
        // TODO 关键字搜索（假设vectorStoreService也支持关键字搜索）
        rerankDocuments = vectorResults.stream().map(Document::getText).toList();


        ReRankRequest request = new ReRankRequest();
        request.setQuery(query);
        request.setDocuments(rerankDocuments);
        request.setTop_n(rerankDocuments.size());
        request.setReturn_documents(false);
        ReRankResponse rerank = openAIService.rerank(request);
        if (rerank == null || rerank.getResults() == null || rerank.getResults().isEmpty()) {
            throw new BaseException("rerank error");
        }
        List<Document> results = new ArrayList<>();
        rerank.getResults().forEach(result -> {
            int index = result.getIndex();
            double relevanceScore = result.getRelevance_score();
            // 根据阈值进行筛选
            if (relevanceScore >= 0.3) {
                // 将筛选后的结果添加到最终结果中
                results.add(vectorResults.get(index));
            }
        });

        return results;
    }

    /**
     * 大模型问答功能
     *
     * @param question 问题
     * @return 回答
     */
    public R<String> answerQuestion(String question) {
        // 使用OpenAI的大模型进行问答
        String answer = openAIService.openAiStream(question).blockFirst();
        return R.success(answer);
    }

    /**
     * 合并并去重搜索结果
     *
     * @param vectorResults 向量搜索结果
     * @param keywordResults 关键字搜索结果
     * @return 合并后的结果
     */
    private List<Document> combineAndDeduplicateResults(List<Document> vectorResults, List<Document> keywordResults) {
        // 合并逻辑，去重等
        // 这里可以根据具体需求实现合并策略
        // 例如，优先保留向量搜索结果，或者根据某种评分机制进行排序
        // 这里简单地将两个结果集合并

        vectorResults.addAll(keywordResults);
        return vectorResults;
    }
}